#ifndef EFFICIENTSAM_H
#define EFFICIENTSAM_H

#include <thread>
#include <mutex>
#include <map>
#include <vector>
#include <opencv2/opencv.hpp>
#include <onnxruntime_cxx_api.h>

class EfficientSam
{
public:
	EfficientSam(const std::string& encoder_path, const std::string& decoder_path);
	void set_image(const cv::Mat& image_input);
	cv::Mat predict_mask_from_points(const std::vector<cv::Point2f>& points, const std::vector<float>& point_labels);
	cv::Mat predict_polygon_from_points(const std::vector<cv::Point2f>& points, const std::vector<float>& point_labels);

private:
	void compute_and_cache_image_embedding(std::size_t image_hash);
	cv::Mat get_image_embedding();

	Ort::Env env;
	Ort::SessionOptions session_options;
	Ort::Session encoder_session;
	Ort::Session decoder_session;

	std::mutex mtx;
	std::map<std::size_t, cv::Mat> image_embedding_cache;
	std::thread worker_thread;

	cv::Mat image;
	cv::Mat image_embedding;
};


// // Usage example
// int main() {
//     EfficientSam sam("encoder.onnx", "decoder.onnx");

//     cv::Mat image = cv::imread("image.jpg");
//     sam.set_image(image);

//     std::vector<cv::Point2f> points = {{100.0f, 150.0f}, {200.0f, 250.0f}};
//     std::vector<float> point_labels = {1.0f, 0.0f};

//     cv::Mat mask = sam.predict_mask_from_points(points, point_labels);

//     cv::imshow("Mask", mask);
//     cv::waitKey(0);

//     return 0;
// }

#endif // EFFICIENTSAM_H
